# sumoVehType2ProbabilityDict = {'passenger':0.8,'bus':0.10,'delivery':0.05,'truck':0.05}
# roufile = '1JuncMorning1.rou.xml'
import xml.etree.ElementTree as ET
import numpy as np
import csv
class RouTools:
    def __init__(self):
        self.csvfile = 'addvehtype.csv'
        self.sumoVehType2ProbabilityDict = {}
        self.readCsv()
    def readCsv(self):
        '''
        读取csv文件
        :return:
        '''
        with open(self.csvfile,'r',encoding='utf-8') as fp:
            reader = csv.reader(fp)
            next(reader)  # 跳过标题
            for row in reader:
                if '' not in row:  # 跳过配置信息不完整的行
                    self.sumoVehType2ProbabilityDict[row[0]] = float(row[1])
    def rouAddVehType(self,rouFile,outFile):
        '''
        向原始sumo车流文件中添加车辆类型
        :param rouFile: 原始车流文件
        :param sumoVehTypeDict: 不同车辆类型对应的出现概率
        :param outFile: 输出车流文件
        :return:
        '''
        sumoVehTypeDict = self.sumoVehType2ProbabilityDict
        sumoVehTypeList = []
        sumoVehTypeChoiceProbabilityList = []
        for sumoVehType,sumoVehTypeChoiceProbability in sumoVehTypeDict.items():
            sumoVehTypeList.append(sumoVehType)
            sumoVehTypeChoiceProbabilityList.append(sumoVehTypeChoiceProbability)
        tree = ET.parse(rouFile)
        root = tree.getroot()
        # root = ET.Element('routes')
        # 添加车辆类型属性
        # 定义车辆类型
        for vehId in self.sumoVehType2ProbabilityDict.keys():
            elem_vtype = ET.Element('vType')
            elem_vtype.set('id', '%s' % (vehId))  # 车辆类型id
            elem_vtype.set('vClass', '%s' % (vehId))  # 车辆类型
            elem_vtype.set('guiShape', '%s' % (vehId))  # 车辆形状
            # elem_vtype.set('tau', '2')  # 车头时距
            # elem_vtype.set('accel', '2.5')  # 启动加速度m/s2
            # elem_vtype.set('decel', '4.5')  # 停车减速度m/s2
            # elem_vtype.set('minGap', '2')  # 车辆间距m
            # elem_vtype.set('maxSpeed', '70')  # 最大速度m/s
            root.insert(0,elem_vtype)
        # 为车辆添加车辆类型属性
        for item in root.findall('vehicle'):
            TypeChoiceIndexList = np.random.choice(a=len(sumoVehTypeChoiceProbabilityList), size=1, replace=False,
                                               p=sumoVehTypeChoiceProbabilityList)
            choiceType = sumoVehTypeList[TypeChoiceIndexList[0]]
            item.set('departLane','best')
            item.set('type', choiceType)
        # ET.dump(root)
        tree.write(outFile)
# rouAddVehType(roufile,sumoVehType2ProbabilityDict,'1JuncMorning2.rou.xml')
if __name__=='__main__':
    main = RouTools()
    inputData = input('需要添加车辆类型的原始文件名：')
    outputData = input('输出的文件名：')
    main.rouAddVehType(inputData,outputData)